#include "ceres/parallel_for.h"
#include <chrono>
#include <cmath>
#include <iostream>

using namespace ceres;
using namespace ceres::internal;
using namespace std;
using namespace std::chrono;

template <typename Scalar>
bool exactMatch(const vector<Scalar>& lst1, const vector<Scalar>& lst2)
{
    if (lst1.size() != lst2.size())
        return false;
    for (size_t i = 0; i < lst1.size(); ++i) {
        if (lst1[i] != lst2[i])
            return false;
    }
    return true;
}

void busyWait(int msec)
{
    auto time_start = system_clock::now();
    while (duration_cast<milliseconds>(system_clock::now() - time_start).count() < msec) {
    }
}

void parallelTest()
{
    ContextImpl context;
    context.EnsureMinimumThreads(/*num_threads=*/2);

    const int size = 100;
    std::vector<int> expected_results(size, 0);
    for (int i = 0; i < size; ++i) {
        expected_results[i] = std::sqrt(i);
    }

    for (int num_threads = 1; num_threads <= 8; ++num_threads) {
        cout << "thread num: " << num_threads << " ---------------------" << endl;
        auto time_start = system_clock::now();
        std::vector<int> values(size, 0);
        ParallelFor(&context, 0, size, num_threads,
            [&values](int i) { values[i] = std::sqrt(i);busyWait(10); });
        cout << "time cost: " << duration_cast<microseconds>(system_clock::now() - time_start).count() / 1000. << " ms" << endl;
        cout << "result: " << (exactMatch(expected_results, values) ? "pass" : "fail") << endl;
    }
}

int main(int, char**)
{
    parallelTest();
    return 0;
}
